#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from retrying import retry
import requests
import json
import fire


@retry
def ElinkSearch(ids: str, start: int, end: int):
    url = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/elink.fcgi'
    payload = {"dbfrom": "protein", "db": "pubmed", "id": ids, "retmode": "json"}
    r = requests.post(url, stream=True, data=payload)
    if r.status_code == 200:
        result = json.loads(r.text)
        pmid_set = set()
        if "linksets" in result:
            for linkset in result["linksets"]:
                if "linksetdbs" in linkset:
                    for linksetdb in linkset["linksetdbs"]:
                        if "links" in linksetdb:
                            for link in linksetdb["links"]:
                                pmid_set.add(link)
        print(str(start) + "~" + str(end) + " added.")
        return (pmid_set)
    else:
        print("Web Error. Retrying ...")
        raise IOError("Web Error.")


def GetPMIDfromProteinID(filepath: str, step: str):
    interval = int(step)
    start = 1
    ids = ""
    total_set = set()
    infile = open(filepath, "r")
    for line in infile:
        line = line.rstrip()
        if line != "":
            start += 1
            ids += line + ","
            if start % interval == 0:
                ids = ids[:-1]
                pmid_set = ElinkSearch(ids, start - 499, start)
                for item in pmid_set:
                    total_set.add(item)
                ids = ""
    infile.close()
    outfile = open("pmids.txt", "w")
    for item in total_set:
        outfile.write(str(item) + "\n")
        outfile.flush()
    outfile.close()


if __name__ == '__main__':
    fire.Fire()
